import os
import torch
import numpy as np
import random
import matplotlib.pyplot as plt
from sklearn.metrics import roc_curve, precision_recall_curve


def set_seed(seed=42):
    """设置随机种子以确保结果可重复"""
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
    os.environ['PYTHONHASHSEED'] = str(seed)


def plot_learning_curves(history, save_path=None):
    """绘制学习曲线"""
    epochs = range(1, len(history['train_loss']) + 1)

    # 创建包含两个子图的图表
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))

    # 绘制损失曲线
    ax1.plot(epochs, history['train_loss'], 'b-', label='训练损失')
    ax1.plot(epochs, history['val_loss'], 'r-', label='验证损失')
    ax1.set_title('训练和验证损失')
    ax1.set_xlabel('轮次')
    ax1.set_ylabel('损失')
    ax1.legend()
    ax1.grid(True)

    # 绘制AUC和准确率曲线
    ax2.plot(epochs, history['val_auc'], 'g-', label='验证AUC')
    ax2.plot(epochs, history['val_acc'], 'p-', label='验证准确率')
    ax2.set_title('验证AUC和准确率')
    ax2.set_xlabel('轮次')
    ax2.set_ylabel('得分')
    ax2.legend()
    ax2.grid(True)

    plt.tight_layout()

    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.close()
    else:
        plt.show()


# ==================== 代码修改部分开始 ====================
def analyze_gate_weights(model, dataloader, device, num_batches=5):
    """分析门控权重的分布 (已修正)"""
    model.eval()
    # 直接针对 'moe_layer' 进行分析
    all_gate_weights = []

    with torch.no_grad():
        for i, batch in enumerate(dataloader):
            if i >= num_batches:
                break

            int_features = batch['int_features'].to(device)
            cat_features = batch['cat_features'].to(device)

            outputs = model(int_features, cat_features)
            gate_weights = outputs['gate_weights']

            # 检查 'moe_layer' 是否存在且不为空
            if 'moe_layer' in gate_weights and gate_weights['moe_layer'] is not None:
                all_gate_weights.append(gate_weights['moe_layer'].cpu().numpy())

    # 如果没有收集到任何权重，则返回None
    if not all_gate_weights:
        return None

    # 计算平均门控权重
    avg_gate_weights = np.mean(np.concatenate(all_gate_weights), axis=0)

    # 返回一个字典，以保持与main.py中循环的兼容性
    return {'moe_layer': avg_gate_weights}
# ==================== 代码修改部分结束 ====================